/*
 * ImmediateCrypt
 * Copyright (C) 2012 Giacomo Drago
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 * 
 * http://giacomodrago.com/go/immediatecrypt
 * 
 */

package com.giacomodrago.immediatecrypt.aes;

import org.apache.commons.io.Charsets;
import org.apache.commons.lang3.RandomStringUtils;
import org.bouncycastle.crypto.BufferedBlockCipher;
import org.bouncycastle.crypto.DataLengthException;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.digests.SHA512Digest;
import org.bouncycastle.crypto.engines.AESEngine;
import org.bouncycastle.crypto.generators.PKCS5S1ParametersGenerator;
import org.bouncycastle.crypto.modes.CBCBlockCipher;
import org.bouncycastle.crypto.paddings.PaddedBufferedBlockCipher;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;

class AESFacadeImpl implements AESFacade {

    private final static int SALT_LENGTH = 16;
    private final static int KEY_SIZE = 256; // AES-256
    private final static int IV_SIZE = 128;
    private final static int PBE_ITERATION_COUNT = 1000;

    @Override
    public AESEncryptedMessage encrypt(byte[] plaintext, String password)
            throws EncryptionException {

        // Check password is not empty
        if (password.isEmpty()) {
            throw new EncryptionException("Password is empty.");
        }
        
        // Generate random password salt
        String salt = RandomStringUtils.randomAlphanumeric(SALT_LENGTH);

        ParametersWithIV params = createEncryptionParameters(password, salt);
        byte[] iv = params.getIV();

        BufferedBlockCipher cipher = new PaddedBufferedBlockCipher(
                new CBCBlockCipher(new AESEngine()));

        cipher.init(true, params);

        byte[] ciphertext = new byte[cipher.getOutputSize(plaintext.length)];

        int outputLen = cipher.processBytes(
                plaintext, 0, plaintext.length, ciphertext, 0);
        try {
            cipher.doFinal(ciphertext, outputLen);
        } catch (DataLengthException ex) {
            throw new EncryptionException(ex);
        } catch (IllegalStateException ex) {
            throw new EncryptionException(ex);
        } catch (InvalidCipherTextException ex) {
            throw new EncryptionException(ex);
        }

        return new AESEncryptedMessage(salt, iv, ciphertext);

    }

    @Override
    public byte[] decrypt(AESEncryptedMessage encryptedMessage, String password)
            throws EncryptionException {

        byte[] ciphertext = encryptedMessage.getCiphertext();
        String salt = encryptedMessage.getSalt();
        byte[] iv = encryptedMessage.getIv();

        ParametersWithIV params = createDecryptionParameters(password, salt, iv);

        BufferedBlockCipher cipher = new PaddedBufferedBlockCipher(
                new CBCBlockCipher(new AESEngine()));

        try {
            cipher.init(false, params);
        } catch (IllegalArgumentException ex) {
            throw new EncryptionException(ex);
        }

        byte[] plaintext = new byte[cipher.getOutputSize(ciphertext.length)];

        int outputLen = cipher.processBytes(
                ciphertext, 0, ciphertext.length, plaintext, 0);
        try {
            cipher.doFinal(plaintext, outputLen);
        } catch (DataLengthException ex) {
            throw new EncryptionException(ex);
        } catch (IllegalStateException ex) {
            throw new EncryptionException(ex);
        } catch (InvalidCipherTextException ex) {
            throw new EncryptionException(ex);
        }

        return plaintext;

    }

    protected ParametersWithIV createDecryptionParameters(String password,
            String salt, byte[] iv) {

        byte[] passwordBytes = password.getBytes(Charsets.UTF_8);
        byte[] saltBytes = salt.getBytes(Charsets.UTF_8);

        PKCS5S1ParametersGenerator keyGenerator = new PKCS5S1ParametersGenerator(
                new SHA512Digest());
        keyGenerator.init(passwordBytes, saltBytes, PBE_ITERATION_COUNT);

        KeyParameter params =
                (KeyParameter) keyGenerator.generateDerivedParameters(KEY_SIZE);

        return new ParametersWithIV(params, iv);

    }

    protected ParametersWithIV createEncryptionParameters(String password, String salt) {

        byte[] passwordBytes = password.getBytes(Charsets.UTF_8);
        byte[] saltBytes = salt.getBytes(Charsets.UTF_8);

        PKCS5S1ParametersGenerator keyGenerator = new PKCS5S1ParametersGenerator(
                new SHA512Digest());
        keyGenerator.init(passwordBytes, saltBytes, PBE_ITERATION_COUNT);

        ParametersWithIV params =
                (ParametersWithIV) keyGenerator.generateDerivedParameters(KEY_SIZE, IV_SIZE);

        return params;

    }
}
